import tensorflow as tf
import matplotlib.pyplot as plt


def data_load(data_dir, img_height, img_width, batch_size):
    train_ds = tf.keras.preprocessing.image_dataset_from_directory(
        data_dir,
        label_mode='categorical',  # 标签变为分类变量
        validation_split=0.2,  # 验证集
        subset='training',
        seed=123,
        image_size=(img_height, img_width),  # 重调大小
        batch_size=batch_size
    )
    val_ds = tf.keras.preprocessing.image_dataset_from_directory(
        data_dir=data_dir,
        label_mode='categorical',  # 标签变为分类变量
        validation_split=0.2,  # 验证集
        subset='validation',
        seed=123,
        image_size=(img_height, img_width),  # 重调大小
        batch_size=batch_size
    )
    class_names = train_ds.classnames
    return train_ds, val_ds, class_names


model = tf.keras.applications.MobileNetV2()
